import numpy as np
from shapely.geometry import Polygon


class DetectionIoUMetrics(object):
    def __init__(self, iou_constraint=0.5):
        self.iou_constraint = iou_constraint

    def evaluate_image(self, gt, pred):
        def get_union(pd, pg):
            return Polygon(pd).union(Polygon(pg)).area

        def get_intersection_over_union(pd, pg):
            return get_intersection(pd, pg) / get_union(pd, pg)

        def get_intersection(pd, pg):
            return Polygon(pd).intersection(Polygon(pg)).area

        det_matched = 0

        pairs = []
        gt_pols = [gt[i] for i in range(len(gt)) if Polygon(gt[i]).is_valid]
        det_pols = [pred[i] for i in range(len(pred)) if Polygon(pred[i]).is_valid]

        if len(gt_pols) > 0 and len(det_pols) > 0:
            output_shape = [len(gt_pols), len(det_pols)]
            iou_mat = np.empty(output_shape)
            gt_rect_mat = np.zeros(len(gt_pols), np.int8)
            det_rect_mat = np.zeros(len(det_pols), np.int8)
            for gt_num in range(len(gt_pols)):
                for det_num in range(len(det_pols)):
                    p_g = gt_pols[gt_num]
                    p_d = det_pols[det_num]
                    iou_mat[gt_num, det_num] = get_intersection_over_union(p_d, p_g)

            for gt_num in range(len(gt_pols)):
                for det_num in range(len(det_pols)):
                    # gt_rect_mat  det_rect_mat用于标志是否已经计算过
                    if gt_rect_mat[gt_num] == 0 and det_rect_mat[det_num] == 0:
                        if iou_mat[gt_num, det_num] > self.iou_constraint:
                            gt_rect_mat[gt_num] = 1
                            det_rect_mat[det_num] = 1
                            det_matched += 1
                            pairs.append({'gt': gt_num, 'det': det_num})

        recall = float(det_matched) / len(gt_pols)
        precision = float(det_matched) / len(det_pols)

        hmean = 0 if (precision + recall) == 0 else 2.0 * precision * recall / (precision + recall)

        per_sample_metrics = {
            'precision': precision,
            'recall': recall,
            'hmean': hmean,
            'pairs': pairs,
            'gtCare': len(gt_pols),
            'detCare': len(det_pols),
            'detMatched': det_matched,
        }
        # import json
        # print(json.dumps(per_sample_metrics, indent=2, ensure_ascii=False))
        return per_sample_metrics

    @staticmethod
    def combine_results(results):
        num_global_care_gt = 0
        num_global_care_det = 0
        matched_sum = 0
        for result in results:
            num_global_care_gt += result['gtCare']
            num_global_care_det += result['detCare']
            matched_sum += result['detMatched']

        method_recall = 0 if num_global_care_gt == 0 else float(
            matched_sum) / num_global_care_gt
        method_precision = 0 if num_global_care_det == 0 else float(
            matched_sum) / num_global_care_det
        method_hmean = 0 if method_recall + method_precision == 0 else 2 * method_recall * method_precision / (
                                                                               method_recall + method_precision)
        method_metrics = {
            'precision': method_precision,
            'recall': method_recall,
            'hmean': method_hmean
        }

        return method_metrics
